# -*- coding: utf-8 -*-

from mtl.op_exec import *
from mtl.op_tran import *
from mtl.op_utils import *
from mtl.hw_defs import VARS

from programs.opcf.mm import *


def conv3d_ld_wait_flag(sm_addr):
    for alu_id in range(VARS.num_of_alu):
        detect_addr = sm_addr + alu_id * VARS.alu_offset
        # Wait the flag.
        clr(pid=mtl_pid_all(), dst=detect_addr, ain=detect_addr)


def conv3d_wino(simd_pid, simd_alu, sm_addr_ia, sm_addr_iaw, sm_addr_kernel,
                sm_addr_tmp2x4, sm_addr_channel_result, sm_addr_channel_flag):
    mul_v(pid=simd_pid, alu=simd_alu,
          dst=sm_addr_iaw,     offset=True, count=16,
          src0=sm_addr_ia,     offset0=True, wait0=True,
          src1=sm_addr_kernel, offset1=False)
    # Set the complete flag.
    add(pid=simd_pid, alu=simd_alu,
        dst=sm_addr_channel_flag, offset=True,
        src0=0, offset0=True,
        src1=1, offset1=True)
    # Do AT P A transform.
    # AT.P
    add_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_tmp2x4,              offset=True,
          src0=sm_addr_iaw,                offset0=True,
          src1=sm_addr_iaw + 4,            offset1=True)
    add_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_tmp2x4,              offset=True,
          src0=sm_addr_tmp2x4,             offset0=True,
          src1=sm_addr_iaw + 8,            offset1=True)
    add_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_tmp2x4 + 4,          offset=True,
          src0=sm_addr_iaw + 4,            offset0=True,
          src1=sm_addr_iaw + 12,           offset1=True)
    sub_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_tmp2x4 + 4,          offset=True,
          src0=sm_addr_tmp2x4 + 4,         offset0=True,
          src1=sm_addr_iaw + 8,            offset1=True)
    # (AT.P).A
    add_v(pid=simd_pid, alu=simd_alu, count=2,
          dst=sm_addr_channel_result,      offset=True, dst_stride=2,
          src0=sm_addr_tmp2x4,             offset0=True, stride0=4,
          src1=sm_addr_tmp2x4 + 1,         offset1=True, stride1=4)
    add_v(pid=simd_pid, alu=simd_alu, count=2,
          dst=sm_addr_channel_result,      offset=True, dst_stride=2,
          src0=sm_addr_channel_result,     offset0=True, stride0=2,
          src1=sm_addr_tmp2x4 + 2,         offset1=True, stride1=4)
    add_v(pid=simd_pid, alu=simd_alu, count=2,
          dst=sm_addr_channel_result + 1,  offset=True, dst_stride=2,
          src0=sm_addr_tmp2x4 + 1,         offset0=True, stride0=4,
          src1=sm_addr_tmp2x4 + 3,         offset1=True, stride1=4)
    sub_v(pid=simd_pid, alu=simd_alu, count=2,
          dst=sm_addr_channel_result + 1,  offset=True, dst_stride=2,
          src0=sm_addr_channel_result + 1, offset0=True, stride0=2,
          src1=sm_addr_tmp2x4 + 2,         offset1=True, stride1=4)


def conv3d_image(simd_pid, simd_alu, sm_addr_ia, sm_addr_iaw, sm_addr_kernel,
                 sm_addr_tmp2x4, sm_addr_channel_result, sm_addr_result,
                 sm_addr_channel_flag):
    # Prepare the channel result.
    channel_addr = sm_addr_channel_result
    # Execute convolution for each channel.
    for kernel_id in range(3):
        conv3d_wino(simd_pid=simd_pid,
                    simd_alu=simd_alu,
                    sm_addr_ia=sm_addr_ia,
                    sm_addr_iaw=sm_addr_iaw,
                    sm_addr_kernel=sm_addr_kernel[kernel_id],
                    sm_addr_tmp2x4=sm_addr_tmp2x4,
                    sm_addr_channel_result=channel_addr,
                    sm_addr_channel_flag=sm_addr_channel_flag)
        channel_addr += 4
    # Sum up the data.
    add_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_result,              offset=True,
          src0=sm_addr_channel_result,     offset0=True,
          src1=sm_addr_channel_result + 4, offset1=True)
    add_v(pid=simd_pid, alu=simd_alu, count=4,
          dst=sm_addr_result,              offset=True,
          src0=sm_addr_result,             offset0=True,
          src1=sm_addr_channel_result + 8, offset1=True)


def conv3d_save(simd_pid, simd_alu, addr_result, sm_addr_result):
    # Send the data back to main memory.
    put_v(dst_pid=mtl_pid_main_mem(), dst=addr_result, dst_inc=2, count=2,
          src_pid=simd_pid, src_alu=simd_alu, src=sm_addr_result)
    put_v(dst_pid=mtl_pid_main_mem(), dst=addr_result + 64, dst_inc=2, count=2,
          src_pid=simd_pid, src_alu=simd_alu, src=sm_addr_result + 2)


def conv3d_64_deploy_kernel(addrs_kernel: list):
    # For 4 ALUs, we just need to deploy to each of them.
    sm_addr_kernel = smalloc('conv3d', 16)
    # Construct the kernel position.
    kernel_addrs = [sm_addr_kernel,
                    sm_addr_kernel + VARS.alu_offset,
                    sm_addr_kernel + VARS.alu_offset * 2]
    # Deploy the data to address.
    ld_v(pid=mtl_pid_all(), dst=kernel_addrs[0], src=addrs_kernel[0], count=16)
    ld_v(pid=mtl_pid_all(), dst=kernel_addrs[1], src=addrs_kernel[1], count=16)
    ld_v(pid=mtl_pid_all(), dst=kernel_addrs[2], src=addrs_kernel[2], count=16)
    # Okay, give back the kernel address.
    return kernel_addrs


def conv3d_64_deploy_rows(sm_addr_ia, addr_channel_ptr, sm_addr_channel_flag):
    for channel_id in range(3):
        for peer_id in range(VARS.num_of_peer):
            for alu_id in range(VARS.num_of_alu):
                ld_v(pid=mtl_pid(peer_id),
                     src=addr_channel_ptr[channel_id],
                     dst=sm_addr_ia + VARS.alu_offset * alu_id,
                     count=16)
                addr_channel_ptr[channel_id] += 16
        # Wait the channel complete flag.
        conv3d_ld_wait_flag(sm_addr_channel_flag)


def conv3d_64(addrs_image, addrs_kernel: list, addr_result):
    # Deploy the kernel to all channel.
    sm_addr_kernel = conv3d_64_deploy_kernel(addrs_kernel)
    # Allocate soft-mem for each alu.
    sm_addr_ia = smalloc('conv3d', 16)
    sm_addr_iaw = smalloc('conv3d', 16)
    sm_addr_tmp2x4 = smalloc('conv3d', 8)
    sm_addr_channel_result = smalloc('conv3d', 4 * 3)
    sm_addr_result = smalloc('conv3d', 4)
    sm_addr_channel_flag = smalloc('conv3d', 1)
    sm_addr_flag = smalloc('conv3d', 1)
    # Deploy rows to chip.
    addr_channel_ptr = addrs_image.copy()
    addr_result_ptr = addr_result

    # Loop for several times.
    for _ in range(16):
        # Deploy the layout.
        conv3d_64_deploy_rows(sm_addr_ia, addr_channel_ptr, sm_addr_channel_flag)
        # Execute the Winograd on all the channels.
        conv3d_image(simd_pid=mtl_pid_all(),
                     simd_alu=mtl_alu_all(),
                     sm_addr_ia=sm_addr_ia,
                     sm_addr_iaw=sm_addr_iaw,
                     sm_addr_kernel=sm_addr_kernel,
                     sm_addr_tmp2x4=sm_addr_tmp2x4,
                     sm_addr_channel_result=sm_addr_channel_result,
                     sm_addr_result=sm_addr_result,
                     sm_addr_channel_flag=sm_addr_channel_flag)
        # Save the data back to main memory.
        conv3d_save(simd_pid=mtl_pid(range(0, 8)),
                    simd_alu=mtl_alu_all(),
                    addr_result=addr_result_ptr,
                    sm_addr_result=sm_addr_result)
        addr_result_ptr += 64 * 2
        conv3d_save(simd_pid=mtl_pid(range(8, 16)),
                    simd_alu=mtl_alu_all(),
                    addr_result=addr_result_ptr,
                    sm_addr_result=sm_addr_result)
        addr_result_ptr += 64 * 2
    # When all the calculation is finished. set the flag of result ptr.
    add(pid=mtl_pid_all(), alu=mtl_alu_all(),
        dst=sm_addr_flag,   offset=True,
        src0=0,             offset0=True,
        src1=1,             offset1=True)
    # Wait for all the ALU complete.
    conv3d_ld_wait_flag(sm_addr_flag)
    # Clear the memory.
    for alu_id in range(VARS.num_of_alu):
        # Wait the flag.
        clr_v(pid=mtl_pid_all(),
              dst=alu_id * VARS.alu_offset,
              count=sm_addr_flag)